#include "main/utils.h"
using namespace std;

int partition(vector<int> &vec, int begin, int end) {
  int value = vec[begin];
  int small = begin;
  for (int i = small + 1; i <= end; ++i) {
    if (vec[i] <= value) {
      ++small;
      if (i != small)
        swap(vec, i, small);
    }
  }
  swap(vec, begin, small);
  return small;
}

vector<int> smallestKNum_partition(vector<int> vec, int k) {
  int begin = 0;
  int end = vec.size() - 1;
  int index = partition(vec, begin, end);
  while (index != k - 1) {
    if (index < k - 1) {
      begin = index + 1;
    } else {
      end = index - 1;
    }
    index = partition(vec, begin, end);
  }
  vector<int> ret(vec.begin(), vec.begin() + k);
  return ret;
}

vector<int> smallestKNum_heap(vector<int> vec, int k) {
  multiset<int, greater<int>> set_tmp;
  for (int i = 0; i < vec.size(); ++i) {
    if (set_tmp.size() < k) {
      set_tmp.insert(vec[i]);
    } else {
      auto top = set_tmp.begin();
      if (vec[i] < *top) {
        set_tmp.erase(top);
        set_tmp.insert(vec[i]);
      }
    }
  }
  vector<int> ret;
  for (auto ele : set_tmp) {
    ret.push_back(ele);
  }
  return ret;
}

int main() {
  vector<int> vec = {4, 5, 1, 6, 2, 7, 3, 8, 0};
  vector<int> ret = smallestKNum_partition(vec, 4);
  printVec("o40_1", ret);

  ret = smallestKNum_heap(vec, 4);
  printVec("o40_2", ret);

  return 0;
}
